package main

import (
	"encoding/xml"
	"fmt"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

type IPSetHeader struct {
	Family     string `xml:"family"`
	Fashsize   int    `xml:"hashsize"`
	Maxelem    int    `xml:"maxelem"`           // 最大成员数
	Timeout    int    `xml:"timeout"`           // 默认超时时间
	Comment    *bool  `xml:"comment,omitempty"` // 是否支持评论 nil 表示不支持
	Memsize    int    `xml:"memsize"`
	References int    `xml:"references"`
	Numentries int    `xml:"numentries"`
}

// 成员
type IPSetMember struct {
	Elem    string `xml:"elem"`    // 值
	Timeout int    `xml:"timeout"` // 超时
	Comment string `xml:"comment"` // 评论
}

type IPSet struct {
	Name     string        `xml:"name,attr"` // 集合名称
	Type     string        `xml:"type"`      // 类型 hash:ip 表示 IP地址
	Revision int           `xml:"revision"`
	Header   IPSetHeader   `xml:"header"`
	Members  []IPSetMember `xml:"members>member"`
}

type Rule struct {
	IP      string `json:"ip" binding:"required"`
	Timeout int    `json:"timeout"` // 超时时间
	Remarks string `json:"remarks"` // 备注
}

type RuleSet struct {
	Name     string `json:"name"`      // 名称
	Timeout  int    `json:"timeout"`   // 超时时间
	MaxRules int    `json:"max_rules"` // 最大规则数量
	Remarks  bool   `json:"remarks"`   // 支持备注
	Rules    []Rule `json:"rules"`     // 规则
	Time     int64  `json:"time"`      // 规则获取时间 unix
}

var rwMutex = sync.RWMutex{}

// 解析ipset xml
func parseIPSetXML(data []byte) (*IPSet, error) {
	var ipsets struct {
		Ipsets []IPSet `xml:"ipset"`
	}
	err := xml.Unmarshal(data, &ipsets)
	if err != nil {
		return nil, err
	}

	if len(ipsets.Ipsets) != 1 {
		return nil, fmt.Errorf("ipset xml error")
	}
	return &ipsets.Ipsets[0], nil
}

// 是否支持评论
func (ipset *IPSet) IsSupportComment() bool {
	return ipset.Header.Comment != nil
}

func (ipset *IPSet) GetRuleSet() *RuleSet {
	rules := make([]Rule, 0)

	for _, member := range ipset.Members {
		rule := Rule{
			IP:      member.Elem,
			Timeout: member.Timeout,
			Remarks: strings.Trim(member.Comment, `"`),
		}
		rules = append(rules, rule)
	}
	set := RuleSet{
		Name:     ipset.Name,
		Timeout:  ipset.Header.Timeout,
		MaxRules: ipset.Header.Maxelem,
		Remarks:  ipset.IsSupportComment(),
		Rules:    rules,
		Time:     time.Now().Unix(),
	}
	return &set
}

// 解析规则集合
func parseRuleSet(data []byte) (*RuleSet, error) {
	ipset, err := parseIPSetXML(data)
	if err != nil {
		return nil, err
	}
	return ipset.GetRuleSet(), nil
}

// 通过ipset集合名称获取xml配置
func getIPSetXML(setname string) ([]byte, error) {
	name := strings.TrimSpace(setname)
	if len(name) == 0 {
		return nil, fmt.Errorf("ipset name cannot be empty")
	}
	cmd := exec.Command("ipset", "list", name, "--output", "xml")
	return cmd.Output()
}

// 通过 ipset 名称获取规则信息
func GetRuleSet(name string) (*RuleSet, error) {
	rwMutex.RLock()
	defer rwMutex.RUnlock()
	raw, err := getIPSetXML(name)
	if err != nil {
		return nil, err
	}
	return parseRuleSet(raw)
}

// 查找字符串中第一个IPV4地址
func MatchIPv4(s string) (string, error) {
	reg, err := regexp.Compile(`(?:25[0-5]|2[0-4]\d|1\d\d|[1-9]\d|\d)(?:\.(?:25[0-5]|2[0-4]\d|1\d\d|[1-9]\d|\d)){3}`)
	if err != nil {
		return "", err
	}
	ip := reg.FindString(strings.TrimSpace(s))
	if ip == "" {
		return "", fmt.Errorf("not match ipv4 address")
	}
	return ip, nil
}

// 实施规则
func (rule *Rule) Apple(setname string) error {
	rwMutex.Lock()
	defer rwMutex.Unlock()
	name := strings.TrimSpace(setname)
	if len(name) == 0 {
		return fmt.Errorf("ipset name cannot be empty")
	}

	ip, err := MatchIPv4(rule.IP)

	if err != nil {
		return err
	}

	remarks := strings.TrimSpace(rule.Remarks)

	reg := regexp.MustCompile("[|\\$>&*;!`\"']")
	if reg.MatchString(remarks) {
		return fmt.Errorf("remarks error")
	}

	// ipset -exist add trust_tmp_ip 116.30.222.169 comment "公司动态IP"
	cmd := exec.Command("ipset", "-exist", "add", name, ip)

	if rule.Timeout > 0 {
		cmd.Args = append(cmd.Args, "timeout", strconv.Itoa(rule.Timeout))
	}

	if len(remarks) != 0 {
		cmd.Args = append(cmd.Args, "comment", remarks)
	}
	err = cmd.Run()

	return err
}

// 删除规则
func (rule *Rule) Remove(setname string) error {
	rwMutex.Lock()
	defer rwMutex.Unlock()
	name := strings.TrimSpace(setname)
	if len(name) == 0 {
		return fmt.Errorf("ipset name cannot be empty")
	}

	ip, err := MatchIPv4(rule.IP)

	if err != nil {
		return err
	}

	// ipset del trust_tmp_ip 116.30.222.169
	cmd := exec.Command("ipset", "-exist", "del", name, ip)

	err = cmd.Run()

	return err
}
